"""
Execute the programs and report execution times. For this script to work,
the wrappers must already be compiled.

This script accepts three parameters : start end job_id
It processes the programs in the range progs_list[start, end)
job_id is used to specify which job generated which list of execution times

The execution times are stored in a list where each entry has the following format :

(function_id, schedule_id, list of execution times)
"""

import sys, pickle, subprocess
from pathlib import Path
from datetime import datetime


start = int(sys.argv[1])
end = int(sys.argv[2])
job_id = int(sys.argv[3])
batchName = sys.argv[4]

# Path to the wrappers (the ones generated by the wrappers job).
data_path = Path("./time_measurement/results_"+batchName)
data_path.mkdir(parents=True, exist_ok=True)

# Path to the list of programs
progs_list_path = Path("./time_measurement/progs_list_"+batchName+".pickle")

# Path to where to store the execution times
dst_path = "./time_measurement/results_"+batchName+"/parts/final_exec_times_%s_%s_%s_{}_{}.pickle" % (job_id, start, end) # This replaces the %s
Path("./time_measurement/results_"+batchName+"/parts").mkdir(parents=True, exist_ok=True)

errors_execute_programs = "./time_measurement/results_"+batchName+"/"+"errors_execute_programs.txt"

# Number of times to execute a single program
nb_tests = 30

# The progress is printed at regular intervals
print_every = 20

# The execution times are saved at regular intervals, in case the script crashes
save_every = 50

with open(progs_list_path, "rb") as f:
    progs_list = pickle.load(f)
    
progs_list = progs_list[start:end]
final_exec_times = []

# Start execution
start = 0
print(str(datetime.now())+"-> %s/%s" % (0, len(progs_list)), flush=True)
for i, prog in enumerate(progs_list):
#     if (i>=3000):
    func_id, sched_id = prog

    # Call run with capture_output=True to get the execution times
    # Execution times are printed to stdout by the wrapper, and capture_output=True uses a pipe for communication
    # With this we can avoid using a disk file for communication
    try:
        proc = subprocess.run([str(data_path / func_id / sched_id / sched_id), str(nb_tests + 1)], check=True, capture_output=True)
        times_list = proc.stdout.decode("utf-8").split(" ")
        times_list = [float(s) for s in times_list[1:-1]]
    except Exception as e:
        times_list = [float(0)]
        with open(errors_execute_programs,"a") as f:
            f.write(sched_id+"     -"+str(job_id)+"-"+str(e)+",\n")
            
    
    

    final_exec_times.append((func_id, sched_id, times_list))

    # Print progress to stdout
    # stdout is redirected to the log file by the command sbatch
    if (i + 1)%print_every == 0:
        print(str(datetime.now())+"-> %s/%s" % (i + 1, len(progs_list)), flush=True)

    # Save the results to disk
    if (i + 1)%save_every == 0:
        with open(dst_path.format(start, i + 1), "wb") as f:
            pickle.dump(final_exec_times, f)

        final_exec_times = []
        start = i + 1
            
#     else:
#         if (i + 1)%print_every == 0:
#             print(str(datetime.now())+"-> %s/%s" % (i + 1, len(progs_list)), flush=True)


# Save remaining results
if final_exec_times != []:
    with open(dst_path.format(start, i + 1), "wb") as f:
        pickle.dump(final_exec_times, f)
